package com.lili.collection.mq.plugins;

import com.lili.collection.core.context.BatchRecordContext;
import com.lili.collection.core.context.StageContext;
import com.lili.collection.core.context.impl.GenericBatchRecordContext;
import com.lili.collection.core.error.InputRunnerException;
import com.lili.collection.core.event.StageEvent;
import com.lili.collection.core.plugins.InputPlugin;
import com.lili.collection.core.runner.StageRunner;
import com.lili.collection.core.runner.result.ErrOutputData;
import com.lili.collection.core.runner.result.RecordInfo;
import com.lili.collection.core.runner.result.RunnerResult;
import com.lili.collection.core.utils.HandlerUtil;
import com.lili.collection.mq.config.PulsarmqInputConfig;
import com.lili.collection.mq.runner.AsyncInputRunner;
import com.lili.collection.mq.utils.MQUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.client.api.*;

import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 描述: kafka消息队列获取数据插件
 *
 * @date:2023/3/17 17:26
 * @author:lili
 */
@Slf4j
public class PulsarmqInputPlugin implements InputPlugin<PulsarmqInputConfig> {

    //别名配置
    private Map<String, String> aliasMap;
    private RecordInfo allRecordInfo;
    private int readNum = 0;
    private PulsarmqInputConfig config;
    private Throwable throwable;
    private Throwable consumerThrowable;
    private PulsarClient client;
    private Consumer<byte[]> clientConsumer;
    private boolean clientClosed;
    private boolean consumerClosed;

    @Override
    public StageContext execute(StageRunner runner) {
        AsyncInputRunner inputRunner = (AsyncInputRunner) runner;
        AtomicBoolean isStart = new AtomicBoolean(false);
        AtomicLong waitTime = new AtomicLong(System.currentTimeMillis());
        try {
            this.clientConsumer = client.newConsumer()
                    .topic(config.getTopic())
                    .subscriptionName(config.getConsumerGroup())
                    .subscriptionType(SubscriptionType.Exclusive)
                    .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
                    .messageListener((consumer, message) -> {
                        try {
                            if(inputRunner.getRunnerTransfer().isRunError()){
                                throw new IllegalStateException("异步运行错误!");
                            }
                            isStart.set(true);
                            this.sendMQDataToBlockQueue(new String(message.getData()),inputRunner);
                            waitTime.set(System.currentTimeMillis());
                            isStart.set(false);
                            if(inputRunner.getRunnerTransfer().isRunError()){
                                throw new IllegalStateException("异步运行错误!");
                            }
                            consumer.acknowledge(message);
                        } catch (Exception e) {
                            isStart.set(false);
                            this.consumerThrowable=e;
                            consumer.negativeAcknowledge(message);
                            dispose();
                        }
                    }).subscribe();
            while(true){
                if(System.currentTimeMillis() - waitTime.get() > 10000 && !isStart.get()){
                    dispose();
                    if(this.consumerThrowable!=null){
                        throw new RuntimeException(this.consumerThrowable);
                    }
                    break;
                }
                if(inputRunner.getRunnerTransfer().isRunError()){
                    throw new IllegalStateException("异步运行错误!");
                }
                TimeUnit.SECONDS.sleep(1);
            }
        } catch (Throwable throwable) {
            runner.getRunnerTransfer().setRunError(true);
            runner.getRunnerTransfer().interruptRunner();
            log.error("获取MQ消息失败!", throwable);
            this.throwable=throwable;
            ErrOutputData errOutputData = allRecordInfo.getErrOutputData();
            errOutputData.setErrorMessage(throwable.toString());
            errOutputData.setType("PULSAR");
        } finally {
            MQUtil.completeStageContext(inputRunner,this.throwable);
        }
        log.info("所有数据读取完成-->读取数据量：{}，异常数：{}，成功数：{}",
                allRecordInfo.getTotalRecordCount(),
                allRecordInfo.getErrorRecordCount(),
                allRecordInfo.getSuccessRecordCount());

        RunnerResult<String> runnerResult = runner.getRunnerResult();
        if(this.throwable==null){
            runnerResult.setTotalRecordCount(allRecordInfo.getTotalRecordCount());
            runnerResult.setSuccessRecordCount(allRecordInfo.getSuccessRecordCount());
            runnerResult.setErrorRecordCount(allRecordInfo.getErrorRecordCount());
        }else {
            runnerResult.setErrOutputData(allRecordInfo.getErrOutputData());
            runnerResult.setSuccessed(false);
        }
        if(inputRunner.getRunnerTransfer()!=null){
            runnerResult.decrementCount(inputRunner.getRunnerTransfer().getDeleteOffsetNum());
        }
        return null;
    }

    private void sendMQDataToBlockQueue(String consumerData, AsyncInputRunner inputRunner) {
        //处理
        if (StringUtils.isBlank(consumerData)) {
            return;
        }
        BatchRecordContext recordContext = new GenericBatchRecordContext();
        RecordInfo tempRecordInfo = HandlerUtil.getRecordInfo(config, consumerData, recordContext);
        if(tempRecordInfo==null){
            return;
        }
        allRecordInfo.addTotalCount(tempRecordInfo.getTotalRecordCount());
        allRecordInfo.addErrorCount(tempRecordInfo.getErrorRecordCount());
        allRecordInfo.addSuccessCount(tempRecordInfo.getSuccessRecordCount());
        if(recordContext.getOutput().size()>0){
            if(inputRunner.getRunnerTransfer().isRunError()){
                throw new IllegalStateException("异步运行错误!");
            }
            inputRunner.getPipeContext().getExchangeStageContext().put(StageRunner.StageType.Input, recordContext);
            inputRunner.triggerStageEvent(StageEvent.STAGE_CONTEXT_CHANGE, inputRunner.getStageType(), recordContext);
            recordContext.setStageContextState(StageContext.StageContextState.CONTINUE);
        }
        log.info("第{}次数据读取完成-->读取数据量：{}，异常数：{}，成功数：{}",
                ++readNum,
                tempRecordInfo.getTotalRecordCount(),
                tempRecordInfo.getErrorRecordCount(),
                tempRecordInfo.getSuccessRecordCount());
    }

    @Override
    public void dispose() {
        try {
            if(this.clientConsumer!=null && !consumerClosed){
                this.clientConsumer.close();
                consumerClosed=true;
            }
        } catch (Exception e) {
            log.error("pulsar消费者关闭异常",e);
        }
        try {
            if (this.client != null && !clientClosed) {
                client.close();
                clientClosed=true;
            }
        } catch (Exception e) {
            log.error("pulsar客户端关闭异常",e);
        }
    }

    @Override
    public void init(PulsarmqInputConfig config) {
        try {
            this.config=config;
            if(this.config==null){
                throw new IllegalArgumentException("pulsar配置数据不能为空!");
            }
            this.config.validate();
            this.aliasMap = config.getAliasMap();
            ClientBuilder clientBuilder = PulsarClient.builder().serviceUrl(config.getServiceUrl());
            clientBuilder.connectionTimeout(10, TimeUnit.SECONDS);
            if(StringUtils.isNotBlank(config.getToken())){
                clientBuilder.authentication(AuthenticationFactory.token(config.getToken()));
            }
            client = clientBuilder.build();
        } catch (Exception e) {
            log.error("创建pulsarmq消费者客户端异常!",e);
            throw new InputRunnerException(e);
        }
        this.allRecordInfo = new RecordInfo();
    }
}